"""
Pedro Martínez Caamaño
Control
Ejercicio 14
"""

import numpy as np
import matplotlib.pyplot as plt


# Apartado a)

def fun_init(N):
    T = np.zeros(N+1)
    T[0] = 0.01
    T[N//2] = 10
    return T

def funcion_calor(N=20, t=15, dt=1e-4, dx=0.1, alpha=1, condiciones_iniciales= fun_init):
    
    r = alpha*(dt/dx**2)
    iteraciones = int(t/dt)
    lista_t = np.arange(0,t, dt)
    
    T = condiciones_iniciales(N)
    Tnew = np.zeros(N+1)
    
    plt.figure(1)
    plt.plot(T)
    
    lista_Tmax = []
    
    for n in range(iteraciones):
        for i in range(1, N):
            if i == 1 or i == N-1:
                Tnew[i] = T[i] + r *(T[i+1] + T[i-1] -2*T[i])
            else:
                Tnew[i] = T[i] + r*(-T[i-2]+16*T[i-1]-30*T[i]+16*T[i+1]-T[i+2])
                
        t = dt * n
        Tnew[-1] = 5*np.sin(t)
        Tnew[0] = T[0] + dt*np.sin(T[0])
        
        T[:] = Tnew[:]
        
        lista_Tmax.append(T.max())
    
        if n%100==0:
            plt.plot(T)

    plt.title("Barra")
    plt.xlabel("i")
    plt.ylabel("T")   
    plt.show()
    return T, np.array(lista_Tmax), lista_t

T, lista_T, lista_t = funcion_calor()

# Apartado b)
plt.figure(2)

plt.plot(lista_t, lista_T, ".")

plt.title("T_max frente a t")
plt.xlabel("t")
plt.ylabel("T_max")   
plt.show()